package com.wtao.testcamera.utils;

import android.media.MediaCodec;
import android.media.MediaCodecInfo;
import android.media.MediaFormat;
import android.media.MediaMuxer;
import android.util.Log;

import androidx.annotation.NonNull;

import java.io.IOException;
import java.nio.ByteBuffer;

public class VideoEncoderAsyncImpl implements IVideoEncoder{
    private static final String TAG = "VideoEncoder";
    private MediaCodec mEncoder;
    private MediaMuxer mMediaMuxer;
    private int mVideoTrackIndex;
    private byte[] mData;
    private boolean isStart = false;
    private boolean mStop = false;
    private static final long DEFAULT_TIMEOUT_US = 50;
    private long prevOutputPTSUs = 0;

    /**
     * 初始化编码器（暂只支持YUV420SP、YUV420P格式）
     * @param outPath 文件输出位置
     * @param width 视频宽度
     * @param height 视频高度
     * @param frameRate 帧率
     * @param bitRate 比特率
     */
    @Override
    public void init(String outPath, int width, int height, int frameRate, int bitRate) {
        try {
            mStop = false;
            isStart = false;
            mVideoTrackIndex = -1;
            mMediaMuxer = new MediaMuxer(outPath, MediaMuxer.OutputFormat.MUXER_OUTPUT_MPEG_4);
            mEncoder = MediaCodec.createEncoderByType(MediaFormat.MIMETYPE_VIDEO_AVC);
            Log.i(TAG, "mMediaMuxer:" + mMediaMuxer + ", mEncoder:" + mEncoder);
            MediaFormat mediaFormat = MediaFormat.createVideoFormat(MediaFormat.MIMETYPE_VIDEO_AVC, width, height);

            mediaFormat.setInteger(MediaFormat.KEY_COLOR_FORMAT, MediaCodecInfo.CodecCapabilities.COLOR_FormatYUV420Flexible);
            mediaFormat.setInteger(MediaFormat.KEY_BIT_RATE, bitRate);
            mediaFormat.setInteger(MediaFormat.KEY_FRAME_RATE, frameRate);
            mediaFormat.setInteger(MediaFormat.KEY_I_FRAME_INTERVAL, 1);

            mEncoder.configure(mediaFormat, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE);
//            mEncoder.start();
            mEncoder.setCallback(new MediaCodec.Callback() {
                @Override
                public void onInputBufferAvailable(@NonNull MediaCodec mediaCodec, int i) {
                    synchronized (this) {
                        try {
                            ByteBuffer inputBuffer = mEncoder.getInputBuffer(i);
                            if(mData == null) {
                                mEncoder.queueInputBuffer(i, 0, 0, 0, 0);
                                return;
                            }
                            inputBuffer.put(mData);
                            mEncoder.queueInputBuffer(i, 0, mData.length, System.nanoTime() / 1000, 0);
                            mData = null;
                        } catch (Exception e) {
                            e.printStackTrace();
                            Log.e(TAG, "!报错！！"+e.getMessage());
                        }
                    }
                }

                @Override
                public void onOutputBufferAvailable(@NonNull MediaCodec mediaCodec, int i, @NonNull MediaCodec.BufferInfo bufferInfo) {
                    try {
                        ByteBuffer outputBuffer = mEncoder.getOutputBuffer(i);
                        if (mVideoTrackIndex == -1) {
                            mVideoTrackIndex = writeHeadInfo(outputBuffer, bufferInfo);
                        }
                        if ((bufferInfo.flags & MediaCodec.BUFFER_FLAG_CODEC_CONFIG) == 0) {
//                        Log.d(TAG, "write outputBuffer");
                            mMediaMuxer.writeSampleData(mVideoTrackIndex, outputBuffer, bufferInfo);
                        }
                        mEncoder.releaseOutputBuffer(i, false);
                    }catch (Exception e) {
                        e.printStackTrace();
                        Log.e(TAG, "报错了!!!！"+e.getMessage());
                    }
                }

                @Override
                public void onError(@NonNull MediaCodec mediaCodec, @NonNull MediaCodec.CodecException e) {
                    Log.e(TAG, "onError: " + e.getMessage());
                }

                @Override
                public void onOutputFormatChanged(@NonNull MediaCodec mediaCodec, @NonNull MediaFormat mediaFormat) {

                }
            });
//            mEncoder.start();
        } catch (IOException e) {
            e.printStackTrace();
            Log.e(TAG, "报错了！"+e.getMessage());
        }
    }

    private int writeHeadInfo(ByteBuffer outputBuffer, MediaCodec.BufferInfo bufferInfo) {
        byte[] csd = new byte[bufferInfo.size];
        outputBuffer.limit(bufferInfo.offset + bufferInfo.size);
        outputBuffer.position(bufferInfo.offset);
        outputBuffer.get(csd);
        ByteBuffer sps = null;
        ByteBuffer pps = null;
        for (int i = bufferInfo.size - 1; i > 3; i--) {
            if (csd[i] == 1 && csd[i - 1] == 0 && csd[i - 2] == 0 && csd[i - 3] == 0) {
                sps = ByteBuffer.allocate(i - 3);
                pps = ByteBuffer.allocate(bufferInfo.size - (i - 3));
                sps.put(csd, 0, i - 3).position(0);
                pps.put(csd, i - 3, bufferInfo.size - (i - 3)).position(0);
            }
        }
        MediaFormat outputFormat = mEncoder.getOutputFormat();
        if (sps != null && pps != null) {
            outputFormat.setByteBuffer("csd-0", sps);
            outputFormat.setByteBuffer("csd-1", pps);
        }
        int videoTrackIndex = mMediaMuxer.addTrack(outputFormat);
        Log.i(TAG, "videoTrackIndex: " + videoTrackIndex);
        mMediaMuxer.start();
        return videoTrackIndex;
    }

    @Override
    public void start() {
        if(mEncoder == null) return;
        if(!isStart) {
            isStart = true;
            mEncoder.start();
        }
    }

    @Override
    public void stop() {
        mEncoder.stop();
        mStop = true;
    }

    @Override
    public void release() {
        if (mEncoder != null) {
            mEncoder.release();
            mEncoder = null;
        }
        if (mMediaMuxer != null) {
            mMediaMuxer.release();
            mMediaMuxer = null;
        }
    }

    @Override
    public void encode(byte[] data, long presentationTimeUs) {
        if (mEncoder == null || mMediaMuxer == null) {
            Log.e(TAG, "mEncoder or mMediaMuxer is null");
            return;
        }
        if (data == null) {
            Log.e(TAG, "input yuv data is null");
            return;
        }
        synchronized (this) {
            mData = data;
        }
    }
}
